"""
WGCNA workflow
"""

from common import execute, load_json
import logging
import re
import os
import csv
import sys
import time
import traceback
import pandas as pd
from argparse import ArgumentParser

PROJECT_DIR = os.path.dirname(os.path.abspath(__file__))
if PROJECT_DIR not in sys.path:
    sys.path.append(PROJECT_DIR)


class WGCNA:
    def __init__(self, param, logger):
        self.config_data = param
        self.start_at = time.perf_counter()
        self.logger = logger
        self.dirs_name = [
            "1_construct_network",
            "2_identify_modules",
            "3_relate_modules_to_traits",
            "4_module_relationships",
            "5_find_key_drivers",
            "6_export_network"
        ]
        self._init_dirs()

    def _write_log(self, log_item: dict):
        func = getattr(self.logger, log_item["type"])
        func(log_item["context"])

    def _init_dirs(self):
        for dir_name in self.dirs_name:
            os.makedirs(
                f"{self.config_data['output']}/{dir_name}", exist_ok=True)
            self._write_log(
                {
                    "type": "info",
                    "context": f"The output directory: {dir_name} was successfully created!",
                }
            )

    def _params_str(self, params: dict) -> str:
        pstr = ""
        for k, v in params.items():
            pstr += f"--{k} {v} "
        return pstr

    def _log_line(self, txt):
        self._write_log(
            {
                "type": "info",
                "context": f"{'-'*20} {txt} {'-'*20}",
            }
        )

    def format_cmd(self, cmds: list) -> list:
        ncmds = []
        for cmd in cmds:
            ncmd = re.sub(r"\s+", " ", cmd).strip()
            ncmds.append(ncmd)
        return ncmds

    def sequential(self, commands):
        for cmd in commands:
            exec_log = execute(cmd)
            self._write_log(exec_log)

    def get_traits(self, filepath):
        with open(filepath,  'r', encoding='utf-8') as file:
            reader = csv.reader(file)
            first_row = next(reader)  # 读取第一行
            first_row_list = list(first_row)  # 转换为列表

        return first_row_list[1:]

    def get_modules(self, filepath):
        df = pd.read_csv(filepath)
        modules = df["Module1"].unique().tolist()
        return [name.replace("ME", "") for name in modules]

    def step1(self):
        step_params = self.config_data["step1"]
        pstr = self._params_str({
            "exprFile": step_params["exprFile"],
            "TOMType": step_params["TOMType"],
            "minModuleSize": step_params["minModuleSize"],
            "maxPower": step_params["maxPower"],
            "mergeCutHeight": step_params["mergeCutHeight"],
            "output": f"{self.config_data['output']}/{self.dirs_name[0]}"
        })
        self._log_line("Step1: Construct Network")
        return f"Rscript {PROJECT_DIR}/steps/step1.R {pstr}"

    def step2(self):
        step_params = self.config_data["step2"]
        pstr = self._params_str({
            "exprFile": self.config_data["step1"]["exprFile"],
            "network": f"{self.config_data['output']}/{self.dirs_name[0]}/network.RData",
            "selectGeneNum": step_params["selectGeneNum"],
            "output": f"{self.config_data['output']}/{self.dirs_name[1]}"
        })
        self._log_line("Step2: Identify Modules")
        return f"Rscript {PROJECT_DIR}/steps/step2.R {pstr}"

    def step3(self):
        step_params = self.config_data["step3"]
        pstr = self._params_str({
            "exprFile": self.config_data["step1"]["exprFile"],
            "traitFile": step_params["traitFile"],
            "network": f"{self.config_data['output']}/{self.dirs_name[0]}/network.RData",
            "output": f"{self.config_data['output']}/{self.dirs_name[2]}"
        })
        self._log_line("Step3: Relate modules to traits")
        return f"Rscript {PROJECT_DIR}/steps/step3.R {pstr}"

    def step4(self):
        pstr = self._params_str({
            "exprFile": self.config_data["step1"]["exprFile"],
            "network": f"{self.config_data['output']}/{self.dirs_name[0]}/network.RData",
            "output": f"{self.config_data['output']}/{self.dirs_name[3]}"
        })
        self._log_line("Step4: Module Relationships")
        return f"Rscript {PROJECT_DIR}/steps/step4.R {pstr}"

    def step5(self, traits):
        cmds = []
        self._log_line("Step5: Find Key Drivers")
        for trait in traits:
            pstr = self._params_str({
                "exprFile": self.config_data["step1"]["exprFile"],
                "network": f"{self.config_data['output']}/{self.dirs_name[0]}/network.RData",
                "traitFile": self.config_data["step3"]["traitFile"],
                "traitName": trait,
                "output": f"{self.config_data['output']}/{self.dirs_name[4]}"
            })

            cmd = f"Rscript {PROJECT_DIR}/steps/step5.R {pstr}"
            cmds.append(cmd)

        return cmds

    def step6(self, modules):
        cmds = []
        self._log_line("Step6: Export Network")
        for moduleName in modules:
            pstr = self._params_str({
                "exprFile": self.config_data["step1"]["exprFile"],
                "network": f"{self.config_data['output']}/{self.dirs_name[0]}/network.RData",
                "TOM": f"{self.config_data['output']}/{self.dirs_name[0]}/TOM.RData",
                "annotFile": self.config_data["step6"]["annotFile"],
                "moduleName": moduleName,
                "topNum": self.config_data["step6"]["topNum"],
                "output": f"{self.config_data['output']}/{self.dirs_name[5]}"
            })

            cmd = f"Rscript {PROJECT_DIR}/steps/step6.R {pstr}"
            cmds.append(cmd)

        return cmds

    def runs(self):
        step1_cmd = self.step1()
        step2_cmd = self.step2()
        step3_cmd = self.step3()
        step4_cmd = self.step4()
        self.sequential(self.format_cmd(
            [step1_cmd, step2_cmd, step3_cmd, step4_cmd]))
        traits = self.get_traits(self.config_data["step3"]["traitFile"])
        step5_cmds = self.step5(traits)
        modules = self.get_modules(
            f"{self.config_data['output']}/{self.dirs_name[3]}/module_relationships.csv")

        step6_cmds = self.step6(modules)
        self.sequential(self.format_cmd(step5_cmds + step6_cmds))


if __name__ == "__main__":
    parser = ArgumentParser(description="")
    parser.add_argument(
        "-c", "--config", type=str, required=True, help=" specify config file "
    )
    args = parser.parse_args()
    config_data = load_json(args.config)

    os.makedirs(config_data['output'], exist_ok=True)

    log_name = int(time.time())
    log_file_path = f"{config_data['output']}/{log_name}.log"
    logging.basicConfig(
        filename=log_file_path,
        filemode="a",
        level=logging.DEBUG,
        format="%(asctime)s - %(name)s - %(filename)s[line:%(lineno)d] - %(levelname)s - %(message)s",
    )
    logging.info("The program begins!")
    try:
        WGCNA(config_data, logging).runs()

    except Exception:
        exc_type, exc_value, exc_traceback_obj = sys.exc_info()
        traceback.print_exception(
            exc_type,
            exc_value,
            exc_traceback_obj,
            file=open(log_file_path, "a", encoding="utf-8"),
        )

        sys.exit(exc_type)

    finally:
        logging.info("The program complete!")
